import os

# --- LangChain Core Imports ---
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# --- LangChain Community Imports ---
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.chat_models import ChatOllama # 使用 Ollama

# --- LangChain HuggingFace Imports ---
from langchain_huggingface import HuggingFaceEmbeddings # 使用 HuggingFace 嵌入

# --- LangChain Text Splitters ---
from langchain.text_splitter import RecursiveCharacterTextSplitter

# --- 配置 ---
PDF_PATH = "example.pdf"  # <--- 重要：替换为你的 PDF 文件路径
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200
EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" # 流行的开源嵌入模型
# EMBEDDING_MODEL_NAME = "BAAI/bge-small-en-v1.5" # 另一个强大的选项
OLLAMA_MODEL_NAME = "gemma3:12b" # <--- 重要：确保你在 Ollama 中拉取了这个模型，或替换为其他已拉取的模型名

# --- 1. 加载 PDF 文档 ---
print(f"从以下路径加载 PDF: {PDF_PATH}")
if not os.path.exists(PDF_PATH):
    print(f"错误: 在 {PDF_PATH} 未找到 PDF 文件")
    exit()

try:
    loader = PyPDFLoader(PDF_PATH)
    pages = loader.load()
    print(f"从 PDF 加载了 {len(pages)} 页。")
    if not pages:
        print("错误: 未从 PDF 加载任何内容。文件是否为空或受密码保护？")
        exit()
except Exception as e:
    print(f"加载 PDF 时出错: {e}")
    exit()


# --- 2. 将文档分割成块 ---
print("将文档分割成块...")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=CHUNK_SIZE,
    chunk_overlap=CHUNK_OVERLAP,
    length_function=len,
    is_separator_regex=False,
)
splits = text_splitter.split_documents(pages)
print(f"分割成 {len(splits)} 个块。")
if not splits:
    print("错误: 文档分割后产生零个块。")
    exit()


# --- 3. 创建嵌入 (使用 HuggingFace 模型) ---
print(f"初始化嵌入模型 ({EMBEDDING_MODEL_NAME})...")
print("这可能需要一些时间来下载模型（如果是第一次运行）。")
try:
    # model_kwargs 指定在哪个设备上运行 ('cpu' 或 'cuda')
    # encode_kwargs 指定是否标准化嵌入 (通常推荐)
    embeddings = HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL_NAME,
        model_kwargs={'device': 'cpu'}, # 如果有兼容的 GPU 和 CUDA，可以改为 'cuda'
        encode_kwargs={'normalize_embeddings': True}
    )
except Exception as e:
    print(f"初始化嵌入模型时出错: {e}")
    print("请确保已安装 'sentence-transformers' 和 'langchain-huggingface'。")
    exit()


# --- 4. 将块存储在向量存储中 (FAISS) ---
# FAISS 是一个高效的内存向量存储
print("创建向量存储 (FAISS 索引)...")
# 优化：检查索引是否已存在，如果存在则加载，避免每次重新计算
FAISS_INDEX_PATH = "my_faiss_index_oss"
if os.path.exists(FAISS_INDEX_PATH):
    print(f"从本地加载 FAISS 索引: {FAISS_INDEX_PATH}")
    try:
         # 注意: 对于较新版本 Langchain 的 load_local 可能需要 allow_dangerous_deserialization=True
         # 请根据你的版本和安全策略决定是否使用
        vectorstore = FAISS.load_local(
            FAISS_INDEX_PATH,
            embeddings,
            allow_dangerous_deserialization=True # 仅当你信任索引来源时使用
        )
        print("向量存储加载成功。")
    except Exception as e:
        print(f"加载 FAISS 索引失败: {e}。将重新创建索引。")
        vectorstore = None
else:
     vectorstore = None

if vectorstore is None:
    try:
        print("正在从文档块创建新的 FAISS 索引...")
        vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
        print("向量存储创建成功。")
        print(f"正在将索引保存到本地: {FAISS_INDEX_PATH}")
        vectorstore.save_local(FAISS_INDEX_PATH)
        print("索引保存成功。")
    except Exception as e:
        print(f"创建向量存储时出错: {e}")
        exit()

# --- 5. 设置 RAG 链 (使用 Ollama LLM) ---
print("设置 RAG 链...")

# 定义检索器
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={'k': 5} # 获取前 5 个最相关的块
)

# 定义 LLM (使用 Ollama)
print(f"初始化 LLM ({OLLAMA_MODEL_NAME} via Ollama)...")
try:
    # 确保 Ollama 服务正在运行!
    llm = ChatOllama(model=OLLAMA_MODEL_NAME, temperature=0.1) # temperature 控制创造性
    # 可选：快速测试 Ollama 连接
    # llm.invoke("Hi there")
    print("Ollama LLM 初始化成功。")
except Exception as e:
    print(f"初始化 Ollama LLM 时出错: {e}")
    print(f"请确保 Ollama 服务正在运行，并且模型 '{OLLAMA_MODEL_NAME}' 已经通过 'ollama pull {OLLAMA_MODEL_NAME}' 拉取。")
    exit()

# 定义提示模板
# 对于某些本地模型，可能需要更明确或简单的提示格式
template = """基于以下上下文信息来回答问题。如果你在上下文中找不到答案，请回答“根据提供的文档，我无法回答这个问题”。

上下文:
{context}

问题:
{question}

回答:"""
prompt = ChatPromptTemplate.from_template(template)

# 辅助函数：格式化检索到的文档
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# 构建 RAG 链
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

print("RAG 链设置完成。可以开始提问。")

# --- 6. 提问 ---
try:
    while True:
        query = input("\n输入你的问题 (或输入 'exit' 退出): ")
        if query.lower() == 'exit':
            break
        if not query:
            continue

        print("\n思考中...")
        # 调用链来获取答案
        answer = rag_chain.invoke(query)

        print("\n回答:")
        print(answer)
        print("-" * 50)

except KeyboardInterrupt:
    print("\n正在退出...")
except Exception as e:
    print(f"\n提问过程中发生错误: {e}")

print("\n应用程序结束。")